/***************************************************************************
 *   Copyright (C) 2010 by programmeboy                                    *
 *   programmeboy@gmail.com                                                *
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 *   This program is distributed in the hope that it will be useful,       *
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
 *   GNU General Public License for more details.                          *
 *                                                                         *
 *   You should have received a copy of the GNU General Public License     *
 *   along with this program; if not, write to the                         *
 *   Free Software Foundation, Inc.                                        *
 *   59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             * 
 ***************************************************************************/

#include <ntifs.h>

#include "common.h"
#include "helper.h"

#define KeGetPcr() ((volatile KPCR * const)__readfsdword(0x1C))

PKDDEBUGGER_DATA64
GetKdDebugerData(
    VOID
)
{
    PDBGKD_GET_VERSION64 KdVersionBlock;
    PKDDEBUGGER_DATA64 KDDebuggerData;
    KdVersionBlock = (PDBGKD_GET_VERSION64)(KeGetPcr()->KdVersionBlock); 
#if 0
    __asm{
        mov eax, fs:[0x1c]  ;get KPCR
        mov eax, [eax+0x34] ;get KdVersionBlock 
        mov KdVersionBlock, eax
    }
#endif

    if (!MmIsAddressValid(KdVersionBlock)){
		return NULL;
	}
	KDDebuggerData = (PKDDEBUGGER_DATA64)(*(PULONG)KdVersionBlock->DebuggerDataList);
	if (!MmIsAddressValid(KDDebuggerData)){
		return NULL;
	}
	if (KDDebuggerData->Header.OwnerTag != KDBG_TAG){
		KdPrint(("(KdGetDebuggerDataBlock) KDDebuggerData failed to get"));
		return NULL;
	}
	return KDDebuggerData;
}

#define PTE_BASE    0xC0000000
#define PDE_BASE    0xC0300000
#define PDE_BASE_PAE 0xc0600000

#define PTI_BITS 10

#define MM_ZERO_PTE 0
#define MM_ZERO_KERNEL_PTE 0

#define MM_ZERO_ACCESS         0  // this value is not used.
#define MM_READONLY            1
#define MM_EXECUTE             2
#define MM_EXECUTE_READ        3
#define MM_READWRITE           4  // bit 2 is set if this is writable.
#define MM_WRITECOPY           5
#define MM_EXECUTE_READWRITE   6
#define MM_EXECUTE_WRITECOPY   7
#define MM_NOCACHE             8

#define PAE_ON (1<<5)

typedef enum {
    VCS_INVALID = 0,
    VCS_VALID,
    VCS_TRANSITION,
    VCS_PAGEDOUT,
    VCS_DEMANDZERO,
    VCS_PROTOTYPE
}VALIDITY_CHECK_STATUS;

typedef struct _MMPTE_SOFTWARE {
    ULONG Valid : 1;
    ULONG PageFileLow : 4;
    ULONG Protection : 5;
    ULONG Prototype : 1;
    ULONG Transition : 1;
    ULONG PageFileHigh : 20;
} MMPTE_SOFTWARE;

typedef struct _MMPTE_TRANSITION {
    ULONG Valid : 1;
    ULONG Write : 1;
    ULONG Owner : 1;
    ULONG WriteThrough : 1;
    ULONG CacheDisable : 1;
    ULONG Protection : 5;
    ULONG Prototype : 1;
    ULONG Transition : 1;
    ULONG PageFrameNumber : 20;
} MMPTE_TRANSITION;

typedef struct _MMPTE_PROTOTYPE {
    ULONG Valid : 1;
    ULONG ProtoAddressLow : 7;
    ULONG ReadOnly : 1;  // if set allow read only access.
    ULONG WhichPool : 1;
    ULONG Prototype : 1;
    ULONG ProtoAddressHigh : 21;
} MMPTE_PROTOTYPE;

typedef struct _MMPTE_HARDWARE {
    ULONG Valid : 1;
    ULONG Write : 1;       // UP version
    ULONG Owner : 1;
    ULONG WriteThrough : 1;
    ULONG CacheDisable : 1;
    ULONG Accessed : 1;
    ULONG Dirty : 1;
    ULONG LargePage : 1;
    ULONG Global : 1;
    ULONG CopyOnWrite : 1; // software field
    ULONG Prototype : 1;   // software field
    ULONG reserved : 1;    // software field
    ULONG PageFrameNumber : 20;
} MMPTE_HARDWARE, *PMMPTE_HARDWARE;

typedef struct _MMPTE {
    union  {
        ULONG Long;
        MMPTE_HARDWARE Hard;
        MMPTE_PROTOTYPE Proto;
        MMPTE_SOFTWARE Soft;
        MMPTE_TRANSITION Trans;
    } u;
} MMPTE, *PMMPTE;

typedef struct _MMPTE_SOFTWARE_PAE {
    ULONGLONG Valid : 1;
    ULONGLONG PageFileLow : 4;
    ULONGLONG Protection : 5;
    ULONGLONG Prototype : 1;
    ULONGLONG Transition : 1;
    ULONGLONG Unused : 20;
    ULONGLONG PageFileHigh : 32;
} MMPTE_SOFTWARE_PAE;

typedef struct _MMPTE_TRANSITION_PAE {
    ULONGLONG Valid : 1;
    ULONGLONG Write : 1;
    ULONGLONG Owner : 1;
    ULONGLONG WriteThrough : 1;
    ULONGLONG CacheDisable : 1;
    ULONGLONG Protection : 5;
    ULONGLONG Prototype : 1;
    ULONGLONG Transition : 1;
    ULONGLONG PageFrameNumber : 24;
    ULONGLONG Unused : 28;
} MMPTE_TRANSITION_PAE;

typedef struct _MMPTE_PROTOTYPE_PAE {
    ULONGLONG Valid : 1;
    ULONGLONG Unused0: 7;
    ULONGLONG ReadOnly : 1;  // if set allow read only access.  LWFIX: remove
    ULONGLONG Unused1: 1;
    ULONGLONG Prototype : 1;
    ULONGLONG Protection : 5;
    ULONGLONG Unused: 16;
    ULONGLONG ProtoAddress: 32;
} MMPTE_PROTOTYPE_PAE;

typedef struct _MMPTE_HARDWARE_PAE {
    ULONGLONG Valid : 1;
    ULONGLONG Write : 1; // UP version
    ULONGLONG Owner : 1;
    ULONGLONG WriteThrough : 1;
    ULONGLONG CacheDisable : 1;
    ULONGLONG Accessed : 1;
    ULONGLONG Dirty : 1;
    ULONGLONG LargePage : 1;
    ULONGLONG Global : 1;
    ULONGLONG CopyOnWrite : 1; // software field
    ULONGLONG Prototype : 1;   // software field
    ULONGLONG reserved0 : 1;  // software field
    ULONGLONG PageFrameNumber : 24;
    ULONGLONG reserved1 : 28;  // software field
} MMPTE_HARDWARE_PAE, *PMMPTE_HARDWARE_PAE;

typedef struct _MMPTE_PAE {
    union  {
        LARGE_INTEGER Long;
        MMPTE_HARDWARE_PAE Hard;
        MMPTE_PROTOTYPE_PAE Proto;
        MMPTE_SOFTWARE_PAE Soft;
        MMPTE_TRANSITION_PAE Trans;
    } u;
} MMPTE_PAE;

typedef MMPTE_PAE *PMMPTE_PAE;

#define MiGetPdeAddress(va)  ((MMPTE*)(((((ULONG)(va)) >> 22) << 2) + PDE_BASE))
#define MiGetPteAddress(va) ((MMPTE*)(((((ULONG)(va)) >> 12) << 2) + PTE_BASE))

#define MiGetPdeAddressPae(va)   ((PMMPTE_PAE)(PDE_BASE_PAE + ((((ULONG)(va)) >> 21) << 3)))
#define MiGetPteAddressPae(va)   ((PMMPTE_PAE)(PTE_BASE + ((((ULONG)(va)) >> 12) << 3)))


#define MmPrint(x)

VALIDITY_CHECK_STATUS
MmIsAddressValidExNotPae(
    IN PVOID Pointer
)
{
    VALIDITY_CHECK_STATUS  Return = VCS_INVALID;
    MMPTE* Pde = MiGetPdeAddress(Pointer);

    if(Pde->u.Hard.Valid){

        MMPTE* Pte = MiGetPteAddress(Pointer);
        if( Pte->u.Hard.Valid ){
            Return = VCS_VALID;
        }else{

            //
            // PTE is not valid
            //

            MMPTE pte = *Pte;

            if(pte.u.Long){
                if( pte.u.Soft.Prototype == 1 ){

                    //MmPrint(("PTE entry is not valid, points to prototype PTE.\n"));
                    // more accurate check should be performed here for pointed prototype PTE!
                    Return = VCS_PROTOTYPE;
                }else{

                    if( pte.u.Soft.Transition != 0 ){

                        //
                        // This is a transition page. Consider it invalid.
                        //

                        //MmPrint(("PTE entry is not valid, points to transition page.\n"));

                        Return = VCS_TRANSITION;
                    }else if (pte.u.Soft.PageFileHigh == 0){

                        //
                        // Demand zero page
                        //

                        Return = VCS_DEMANDZERO;
                    }else{

                        //
                        // Pagefile PTE
                        //

                        if( pte.u.Soft.Transition == 0 ){
                            Return = VCS_PAGEDOUT;
                        }
                    }
                }
            }
        }
    }

    return Return;
}


VALIDITY_CHECK_STATUS
MmIsAddressValidExPae(
    IN PVOID Pointer
)
{
    VALIDITY_CHECK_STATUS Return = VCS_INVALID;

    MMPTE_PAE* Pde = MiGetPdeAddressPae(Pointer);

    if( Pde->u.Hard.Valid )
    {
        MMPTE_PAE* Pte;
        if( Pde->u.Hard.LargePage != 0 ){
            Pte = Pde;
        }else{

            //
            // Small 4K page
            //

            Pte  = MiGetPteAddressPae(Pointer);
        }

        if( Pte->u.Hard.Valid ){
            Return = VCS_VALID;
        }else{

            //
            // PTE is not valid
            //

            MMPTE_PAE pte = *Pte;

            if( pte.u.Long.LowPart == 0 ){
                //MmPrint(("PTE entry is completely invalid (page is not committed or is within VAD tree)\n"));
            }else{
                if( pte.u.Soft.Prototype == 1 ){
                
                    //
                    // more accurate check should be performed here for pointed prototype PTE!
                    //

                    Return = VCS_PROTOTYPE;
                }else{

                    //
                    // not a prototype PTE
                    //

                    if( pte.u.Soft.Transition != 0 ){

                        //
                        // This is a transition page.
                        //

                        Return = VCS_TRANSITION;
                    }else if (pte.u.Soft.PageFileHigh == 0){
                        //
                        // Demand zero page
                        //

                        Return = VCS_DEMANDZERO;
                    }else{
                        //
                        // Pagefile PTE
                        //

                        if( pte.u.Soft.Transition == 0 ){

                            Return = VCS_PAGEDOUT;
                        }
                    }
                }
            }
        }
    }

    return Return;
}

VALIDITY_CHECK_STATUS
MmIsAddressValidEx(
    IN PVOID Pointer
)
{
    if((ULONG)__readcr4() & PAE_ON) {
        return MmIsAddressValidExPae(Pointer);
    }else {
        return MmIsAddressValidExNotPae(Pointer);
    }
}


BOOLEAN
MmIsAddressRangeValid(
    IN PVOID Address,
    IN ULONG Size
)
{
    ULONG Ptr = NULL;

    Address = ALIGN_DOWN_POINTER(Address, PAGE_SIZE);
    Size = ALIGN_UP(Size, PAGE_SIZE);

    for( Ptr = (ULONG)Address; 
         Ptr < (ULONG)Address + Size; 
         *(ULONG*)&Ptr += PAGE_SIZE ){
        if( MmIsAddressValidEx((PVOID)Ptr) == VCS_INVALID ) {
            return FALSE;
        }
    }

    return TRUE;
}



